import fitz
import base64
import os
from dotenv import load_dotenv
import requests
from langchain_core.documents import Document
import urllib3

# 禁用SSL警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# 环境变量加载
load_dotenv()

output_dir = "temp_images"
ZHIPU_API_KEY = os.getenv("ZHIPUAI_API_KEY")
pdf_path = r"D:\ideaSpace\rag-in-action-master\90-文档-Data\黑悟空\黑神话悟空.pdf"

def analyze_image_with_glm4(image_path):
    """使用智谱GLM-4模型分析图片（修正版）"""
    with open(image_path, "rb") as image_file:
        base64_image = base64.b64encode(image_file.read()).decode('utf-8')

    headers = {
        "Authorization": f"Bearer {ZHIPU_API_KEY}",
        "Content-Type": "application/json"
    }

    # 修正后的payload结构
    payload = {
        "model": "glm-4v",  # 使用支持多模态的模型
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        # "text": "请详细描述这张PPT幻灯片的内容，包括标题、正文和图片内容。"
                        "text": "请详细描述这张图片的内容"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"  # 修正为jpeg格式
                        }
                    }
                ]
            }
        ],
        "max_tokens": 2000
    }

    try:
        response = requests.post(
            "https://open.bigmodel.cn/api/paas/v4/chat/completions",
            headers=headers,
            json=payload,
            verify=False
        )
        response.raise_for_status()  # 检查HTTP错误
        return response.json()["choices"][0]["message"]["content"]
    except requests.exceptions.HTTPError as http_err:
        print(f"HTTP错误: {http_err}")
        print(f"响应内容: {response.text}")
    except Exception as e:
        print(f"其他错误: {e}")
    return "图片分析失败"

# PDF转图片函数保持不变
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

print("正在转换PDF为图片...")
try:
    doc = fitz.open(pdf_path)
    image_paths = []
    for i in range(len(doc)):
        page = doc.load_page(i)
        pix = page.get_pixmap(dpi=200)
        image_path = os.path.join(output_dir, f"page_{i+1}.png")
        pix.save(image_path)
        image_paths.append(image_path)
    print(f"成功转换 {len(image_paths)} 页")
except Exception as e:
    print(f"PDF转换失败: {e}")
    exit()

# 分析图片
print("\n开始分析图片...")
results = []
for i, image_path in enumerate(image_paths[:2]):  # 先测试前2页
    print(f"正在分析第 {i+1}/{len(image_paths)} 页...")
    try:
        analysis_result = analyze_image_with_glm4(image_path)
        results.append(analysis_result)
        print(f"分析结果: {analysis_result[:100]}...")  # 打印部分结果
    except Exception as e:
        print(f"分析第 {i+1} 页失败: {str(e)}")
        results.append("分析失败")

# 转换为Document格式
documents = [
    Document(
        page_content=result,
        metadata={
            "source": pdf_path,
            "page_number": i + 1,
            "analyzer": "GLM-4V"
        }
    )
    for i, result in enumerate(results)
]

# 输出结果
print("\n分析结果：")
for doc in documents:
    print(f"页码: {doc.metadata['page_number']}")
    print(f"内容: {doc.page_content}")
    print("-" * 80)

# 清理临时文件
print("\n清理临时文件...")
for image_path in image_paths:
    if os.path.exists(image_path):
        os.remove(image_path)
if os.path.exists(output_dir):
    os.rmdir(output_dir)
print("处理完成！")